from io import BytesIO
from typing import List
from urllib.parse import quote

import pandas as pd
import xlwt
from flask import make_response

from app.common.base_manager import BaseManager


class ExcelError(Exception):
    pass


class ExcelManager(BaseManager):
    """
    excel管理器
    since: 22.2.24
    """

    Error = ExcelError

    @staticmethod
    def make_excel_response_by_wbk(wbk, file_name=None):
        """
        生成excel的response
        """
        sio = BytesIO()
        wbk.save(sio)
        # 校验文件名
        file_name = file_name or '导入模板'
        if not file_name.endswith(".xls"):
            file_name += ".xls"
        file_name = quote(file_name)
        response = make_response(sio.getvalue())
        response.headers[
            "Content-Type"
        ] = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
        response.headers["Cache-Control"] = "no-cache"
        response.headers[
            "Content-Disposition"
        ] = f"attachment; filename={file_name}; filename*=utf-8''{file_name}"
        sio.seek(0)
        sio.close()
        return response

    @staticmethod
    def make_wbk_or_add_sheet(sheet_name, heads: List[str] = None, wbk=None):
        """
        创建workbook对象 及添加sheet 加入基本的head结构
        """
        wbk = wbk or xlwt.Workbook(encoding="utf-8")
        sheet = wbk.add_sheet(sheet_name)
        heads = heads or []
        for i, v in enumerate(heads):
            sheet.write(0, i, v)
        return wbk

    def import_template_download(self, file_name: str, sheet_name: str, heads: List[str] = None):
        """
        导入模板下载
        :param file_name: 文件名称
        :param sheet_name: sheet 名称
        :param heads: 表头
        """
        wbk = self.make_wbk_or_add_sheet(sheet_name, heads=heads)
        return self.make_excel_response_by_wbk(wbk, file_name)

    def verify_excel(
        self,
        file,
        verify_heads: List[str] = None,
        notnull_heads: List[str] = None,
    ):
        """
        检验excel文件
        :param file: 文件/文件名称
        :param verify_heads: 校验excel头
        :param notnull_heads: 校验excel非空
        """

        if not file:
            self.error_func("文件不存在!")

        df = None
        try:
            # 注意这里将表格数据都转为 str 类型
            df = pd.read_excel(file.read(), dtype=str)
        except Exception as e:
            print(e)
            self.error_func("文件格式不正确!")

        if all([verify_heads, df.columns.values.tolist() != verify_heads]):
            self.error_func("Excel表格头不正确, 请使用正确的模板!")

        if notnull_heads:
            notnull_frame = pd.DataFrame(df, columns=notnull_heads)
            if pd.isnull(notnull_frame).values.any():
                self.error_func("请完善必填项!")

        # 空表格补空字符串
        df.fillna('', inplace=True)

        return df


excel_manager = ExcelManager()
